{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# FFT\n",
    "\n",
    "Code examples from [Think Complexity, 2nd edition](http://greenteapress.com/wp/complexity2), Appendix A\n",
    "\n",
    "Copyright 2017 Allen Downey, [MIT License](http://opensource.org/licenses/MIT)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from __future__ import print_function, division\n",
    "\n",
    "%matplotlib inline\n",
    "\n",
    "import os\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "import thinkplot\n",
    "\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Empirical order of growth\n",
    "\n",
    "Sometimes we can figure out what order of growth a function belongs to by running it with a range of problem sizes and measuring the run time.\n",
    "\n",
    "To measure runtimes, we'll use `etime`, which uses `os.times` to compute the total time used by a process, including \"user time\" and \"system time\".  User time is time spent running your code; system time is time spent running operating system code on your behalf."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def etime():\n",
    "    \"\"\"Measures user and system time this process has used.\n",
    "\n",
    "    Returns the sum of user and system time.\"\"\"\n",
    "    user, sys, chuser, chsys, real = os.times()\n",
    "    return user+sys"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`time_func` takes a function object and a problem size, `n`, runs the function, and returns the elapsed time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def time_func(func, n):\n",
    "    \"\"\"Run a function and return the elapsed time.\n",
    "    \n",
    "    func: function\n",
    "    n: problem size\n",
    "    \n",
    "    returns: user+sys time in seconds\n",
    "    \"\"\"\n",
    "    start = etime()\n",
    "    func(n)\n",
    "    end = etime()\n",
    "    elapsed = end - start\n",
    "    return elapsed"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`run_timing_test` takes a function, runs it with a range of problem sizes, and returns two lists: problem sizes and times."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def run_timing_test(func, max_time=1):\n",
    "    \"\"\"Tests the given function with a range of values for n.\n",
    "    \n",
    "    func: function object\n",
    "\n",
    "    returns: list of ns and a list of run times.\n",
    "    \"\"\"\n",
    "    ns = []\n",
    "    ts = []\n",
    "    for i in range(6, 28):\n",
    "        n = 2**i\n",
    "        t = time_func(func, n)\n",
    "        print(n, t)\n",
    "        if t > 0:\n",
    "            ns.append(n)\n",
    "            ts.append(t)\n",
    "        if t > max_time:\n",
    "            break\n",
    "\n",
    "    return ns, ts"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`fit` takes the lists of ns and ts and fits it with a curve of the form `a * n**exp`, where `exp` is a given exponent and `a` is chosen so that the line goes through a particular point in the sequence, usually the last. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def fit(ns, ts, exp=1.0, index=-1):\n",
    "    \"\"\"Fits a curve with the given exponent.\n",
    "    \n",
    "    ns: sequence of problem sizes\n",
    "    ts: sequence of times\n",
    "    exp: exponent of the fitted curve\n",
    "    index: index of the element the fitted line should go through\n",
    "    \n",
    "    returns: sequence of fitted times\n",
    "\n",
    "    \n",
    "    \"\"\"\n",
    "    # Use the element with the given index as a reference point, \n",
    "    # and scale all other points accordingly.\n",
    "    nref = ns[index]\n",
    "    tref = ts[index]\n",
    "\n",
    "    tfit = []\n",
    "    for n in ns:\n",
    "        ratio = n / nref\n",
    "        t = ratio**exp * tref\n",
    "        tfit.append(t)\n",
    "\n",
    "    return tfit"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`plot_timing_test` plots the results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def plot_timing_test(ns, ts, label='', color='blue', exp=1.0, scale='log'):\n",
    "    \"\"\"Plots data and a fitted curve.\n",
    "\n",
    "    ns: sequence of n (problem size)\n",
    "    ts: sequence of t (run time)\n",
    "    label: string label for the data curve\n",
    "    color: string color for the data curve\n",
    "    exp: exponent (slope) for the fitted curve\n",
    "    \"\"\"\n",
    "    tfit = fit(ns, ts, exp)\n",
    "    plt.plot(ns, tfit, color='0.7', linewidth=2, linestyle='dashed')\n",
    "    plt.plot(ns, ts, 's-', label=label, color=color, alpha=0.5, linewidth=3)\n",
    "    plt.xlabel('Problem size (n)')\n",
    "    plt.ylabel('Runtime (seconds)')\n",
    "    plt.xscale(scale)\n",
    "    plt.yscale(scale)\n",
    "    plt.legend()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For small values of `n`, the runtime is so short that we're probably not getting an accurate measurement of just the operation we're interested in.  But as `n` increases, runtime seems to converge to a line with slope 1.  \n",
    "\n",
    "That suggests that performing append `n` times is linear, which suggests that a single append is constant time.  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### DFT\n",
    "\n",
    "Here's an implementation of DFT using outer product to construct the transform matrix, and dot product to compute the DFT."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "PI2 = 2 * np.pi\n",
    "\n",
    "def dft(xs):\n",
    "    N = len(xs)\n",
    "    ns = np.arange(N) / N\n",
    "    ks = np.arange(N)\n",
    "    args = np.outer(ks, ns)\n",
    "    M = np.exp(-1j * PI2 * args)\n",
    "    amps = M.dot(xs)\n",
    "    return amps"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here's an example comparing this implementation of DFT with `np.fft.fft`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "xs = np.random.normal(size=128)\n",
    "spectrum1 = np.fft.fft(xs)\n",
    "plt.plot(np.abs(spectrum1))\n",
    "\n",
    "spectrum2 = dft(xs)\n",
    "plt.plot(np.abs(spectrum2))\n",
    "\n",
    "np.allclose(spectrum1, spectrum2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, let's see what the asymptotic behavior of `np.fft.fft` looks like:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_fft(n):\n",
    "    xs = np.random.normal(size=n)\n",
    "    spectrum = np.fft.fft(xs)\n",
    "\n",
    "ns, ts = run_timing_test(test_fft)\n",
    "plot_timing_test(ns, ts, 'test_fft', exp=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Up through the biggest array I can handle on my computer, it's hard to distinguish from linear.\n",
    "\n",
    "And let's see what my implementation of DFT looks like: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_dft(n):\n",
    "    xs = np.random.normal(size=n)\n",
    "    spectrum = dft(xs)\n",
    "\n",
    "ns, ts = run_timing_test(test_dft)\n",
    "plot_timing_test(ns, ts, 'test_dft', exp=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It's definitely much slower than `np.fft.fft`.\n",
    "\n",
    "But it looks like it's not as bad as quadratic.  Maybe it's linear?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_timing_test(ns, ts, 'test_dft', exp=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "No, it looks steeper than linear.  Just based on this data, we would have a hard time classifying this algorithm."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Implementing FFT\n",
    "\n",
    "Ok, let's try our own implementation of FFT.\n",
    "\n",
    "First I'll get the divide and conquer part of the algorithm working:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def fft_norec(ys):\n",
    "    N = len(ys)\n",
    "    He = dft(ys[::2])\n",
    "    Ho = dft(ys[1::2])\n",
    "    \n",
    "    ns = np.arange(N)\n",
    "    W = np.exp(-1j * PI2 * ns / N)\n",
    "    \n",
    "    return np.tile(He, 2) + W * np.tile(Ho, 2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This version breaks the array in half, uses `dft` to compute the DFTs of the halves, and then uses the D-L lemma to stich the results back up.\n",
    "\n",
    "Let's see what the performance looks like."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_fft_norec(n):\n",
    "    xs = np.random.normal(size=n)\n",
    "    spectrum = fft_norec(xs)\n",
    "\n",
    "ns, ts = run_timing_test(test_fft_norec)\n",
    "plot_timing_test(ns, ts, 'test_fft_norec', exp=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Looks about the same as DFT.  Again, it seems better than quadratic, but not as good as linear."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_timing_test(ns, ts, 'test_fft_norec', exp=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Exercise:**  Starting with fft_norec, write a function called fft_rec that's fully recursive; that is, instead of using `dft` to compute the DFTs of the halves, it should use `fft_rec`.  Of course, you will need a base case to avoid an infinite recursion.  You have two options:\n",
    "\n",
    "1) The DFT of an array with length 1 is the array itself.\n",
    "\n",
    "2) If the parameter, `ys`, is smaller than some threshold length, you could use DFT.\n",
    "\n",
    "Use `test_fft_rec`, below, to check the performance of your function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Solution goes here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "xs = np.random.normal(size=128)\n",
    "spectrum1 = np.fft.fft(xs)\n",
    "\n",
    "spectrum2 = fft_rec(xs)\n",
    "\n",
    "np.allclose(spectrum1, spectrum2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_fft_rec(n):\n",
    "    xs = np.random.normal(size=n)\n",
    "    spectrum = fft_rec(xs)\n",
    "\n",
    "ns, ts = run_timing_test(test_fft_rec)\n",
    "plot_timing_test(ns, ts, 'test_fft_rec', exp=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If things go according to plan, your FFT implementation should be faster than `dft` and `fft_norec`, but probably not as fast as `np.fft.fft`.  And it will probably be a bit steeper than linear."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
